package com.farm.controller;

import com.farm.chat.StreamChat;
import lombok.extern.slf4j.Slf4j;
import okhttp3.ResponseBody;
import org.json.JSONException;
import org.json.JSONObject;
import org.springframework.stereotype.Component;

import javax.websocket.*;
import javax.websocket.server.PathParam;
import javax.websocket.server.ServerEndpoint;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CopyOnWriteArraySet;


@Slf4j
@Component
@ServerEndpoint("/websocket/{target}") //创建ws的请求路径。
public class WebsocketServerEndpoint {
    private Session session;
    private String target;
    //支持持续流推送
    private InputStream inputStream;
    private final static CopyOnWriteArraySet<WebsocketServerEndpoint> websockets = new CopyOnWriteArraySet<>();

    @OnOpen
    public void onOpen(Session session , @PathParam("target") String target){
        this.session = session;
        this.target = target;
        websockets.add(this);
        log.info("websocket connect server success , target is {},total is {}",target,websockets.size());
    }

    @OnMessage
    public void onMessage(String message) throws IOException, JSONException {
        log.info("message is {}",message);
        JSONObject jsonObject = new JSONObject(message);
        String user = (String)jsonObject.get("user");
        String question = (String)jsonObject.get("message");

        StreamChat streamChat = new StreamChat();
        ResponseBody body = streamChat.getAnswerStream(question);
        InputStream inputStream = body.byteStream();

        sendMessageSync(user,inputStream);
    }

    @OnClose
    public void onClose(){
        log.info("connection has been closed ,target is {},total is {}" ,this.target, websockets.size());
        this.destroy();
    }

    @OnError
    public void onError(Throwable throwable){
        this.destroy();
        log.info("websocket connect error , target is {} ,total is {}, error is {}",this.target ,websockets.size(),throwable.getMessage());
    }

    /**
     * 根据目标身份推送消息
     * @param target
     * @param message
     * @throws IOException
     */
    public void sendMessageOnce(String target, String message) throws IOException {
        this.sendMessage(target,message,false,null);
    }

    /**
     * stream 同步日志输出，通过websocket推送至前台。
     * @param target
     * @param is
     * @throws IOException
     */
    private void sendMessageSync(String target, InputStream is) throws IOException {
        WebsocketServerEndpoint websocket = getWebsocket(target);
        if (Objects.isNull(websocket)) {
            throw new RuntimeException("The websocket does not exist or has been closed.");
        }
        if (Objects.isNull(is)) {
            throw new RuntimeException("InputStream cannot be null.");
        } else {
            websocket.inputStream = is;
            CompletableFuture.runAsync(websocket::sendMessageWithInputSteam);
        }
    }


    /**
     * Send message.
     * @param target 通过target获取{@link WebsocketServerEndpoint}.
     * @param message message
     * @param continuous 是否通过inputStream持续推送消息。
     * @param is 输入流
     * @throws IOException
     */
    private void sendMessage(String target , String message ,Boolean continuous , InputStream is) throws IOException {
        WebsocketServerEndpoint websocket = getWebsocket(target);
        if(Objects.isNull(websocket)){
            throw new RuntimeException("The websocket does not exists or has been closed.");
        }
        if(continuous){
            if(Objects.isNull(is)){
                throw new RuntimeException("InputStream can not be null when continuous is true.");
            }else{
                websocket.inputStream = is;
                CompletableFuture.runAsync(websocket::sendMessageWithInputSteam);
            }
        }else{
            websocket.session.getBasicRemote().sendText(message);
        }
    }

    /**
     * 通过inputStream 持续推送消息。
     * 支持文件、消息、日志等。
     */
    private void sendMessageWithInputSteam() {
        String message;
        BufferedReader bufferedReader = new BufferedReader(new InputStreamReader(this.inputStream));
        try {
            while ((message = bufferedReader.readLine()) != null) {
                if(message.equals(""))
                    continue;
                if (websockets.contains(this)) {
                    System.out.println(message);
                    this.session.getBasicRemote().sendText(message);
                }
            }
        } catch (IOException e) {
            log.warn("SendMessage failed {}", e.getMessage());
        } finally {
            this.closeInputStream();
        }
    }

    /**
     * 根据目标获取对应的{@link WebsocketServerEndpoint}。
     * @param target 约定标的
     * @return WebsocketServerEndpoint
     */
    private WebsocketServerEndpoint getWebsocket(String target){
        WebsocketServerEndpoint websocket = null;
        for (WebsocketServerEndpoint ws : websockets) {
            if (target.equals(ws.target)) {
                websocket = ws;
            }
        }
        return websocket;
    }

    /**
     * close inputStream.
     * @Author Doit
     * @Date 20221/08/23 15:30:00
     */
    private void closeInputStream(){
        if(Objects.nonNull(inputStream)){
            try {
                inputStream.close();
            } catch (Exception e) {
                log.warn("websocket close failed {}",e.getMessage());
            }
        }
    }

    /**
     * destroy {@link WebsocketServerEndpoint}
     * @Author Doit
     * @Date 20221/08/23 15:30:00
     */
    private void destroy(){
        websockets.remove(this);
        this.closeInputStream();
    }
}
